// agent.go (最终修复版)
package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	// "time" 已被移除

	"github.com/cilium/ebpf/link"
	"github.com/cilium/ebpf/ringbuf"
	"github.com/cilium/ebpf/rlimit"
	"golang.org/x/sys/unix" // <--- 新增的 import
)

//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -target ${GOARCH} bpf ./ebpf/prefetch_kern.c -- -O2 -g -Wall -I./ebpf

const (
	// 预取深度：预测未来多少个块
	PrefetchDepth = 4
	// 页大小，通常是 4KB
	PageSize = 4096
)

// 从 eBPF C 代码中镜像的事件结构
type Event struct {
	Pid       uint32
	FaultAddr uint64
	Stride    int64
}

// 缓存 /proc/[pid]/maps 的解析结果
type VMAEntry struct {
	StartAddr uint64
	EndAddr   uint64
	Offset    uint64
	Path      string
}

var vmaCache = make(map[uint32][]VMAEntry)
var fileCache = make(map[string]*os.File)

func main() {
	stopper := make(chan os.Signal, 1)
	signal.Notify(stopper, os.Interrupt, syscall.SIGTERM)

	if err := rlimit.RemoveMemlock(); err != nil {
		log.Fatal(err)
	}

	objs := bpfObjects{}
	if err := loadBpfObjects(&objs, nil); err != nil {
		log.Fatalf("loading objects: %v", err)
	}
	defer objs.Close()

	tp, err := link.Tracepoint("exceptions", "page_fault_user", objs.HandlePageFault, nil)
	if err != nil {
		log.Fatalf("attaching tracepoint: %v", err)
	}
	defer tp.Close()

	log.Println("eBPF program attached. Waiting for page fault events...")

	rd, err := ringbuf.NewReader(objs.Rb)
	if err != nil {
		log.Fatalf("opening ringbuf reader: %v", err)
	}
	defer rd.Close()

	go func() {
		<-stopper
		log.Println("Received signal, exiting...")
		rd.Close()
	}()

	var event Event
	for {
		record, err := rd.Read()
		if err != nil {
			if errors.Is(err, ringbuf.ErrClosed) {
				log.Println("Ringbuf closed, exiting.")
				return
			}
			log.Printf("reading from reader: %v", err)
			continue
		}

		if err := binary.Read(bytes.NewBuffer(record.RawSample), binary.LittleEndian, &event); err != nil {
			log.Printf("parsing ringbuf event: %v", err)
			continue
		}

		log.Printf("Detected stable pattern for PID %d: Addr=0x%x, Stride=%d bytes", event.Pid, event.FaultAddr, event.Stride)
		
		go handlePrefetch(event)
	}
}

func handlePrefetch(event Event) {
	for i := 1; i <= PrefetchDepth; i++ {
		prefetchAddr := event.FaultAddr + uint64(i)*uint64(event.Stride)

		fd, offset, err := translateVaddr(event.Pid, prefetchAddr)
		if err != nil {
			continue
		}
		if fd == nil {
			continue
		}

		err = doPrefetch(fd, int64(offset), PageSize)
		if err != nil {
			continue
		}
	}
}

// --- 这是关键修改 ---
func doPrefetch(fd *os.File, offset int64, length int) error {
	// 在 Go 的 unix 包中，正确的常量名是 FADV_WILLNEED
	err := unix.Fadvise(int(fd.Fd()), offset, int64(length), unix.FADV_WILLNEED)
	return err
}

func translateVaddr(pid uint32, vaddr uint64) (*os.File, uint64, error) {
	vmas, ok := vmaCache[pid]
	if !ok {
		var err error
		vmas, err = parseProcMaps(pid)
		if err != nil {
			return nil, 0, err
		}
		vmaCache[pid] = vmas
	}

	for _, vma := range vmas {
		if vaddr >= vma.StartAddr && vaddr < vma.EndAddr {
			fileOffset := vma.Offset + (vaddr - vma.StartAddr)

			file, ok := fileCache[vma.Path]
			if !ok {
				if vma.Path == "" || !strings.HasPrefix(vma.Path, "/") {
					return nil, 0, nil
				}
				var err error
				file, err = os.OpenFile(vma.Path, os.O_RDONLY, 0)
				if err != nil {
					return nil, 0, err
				}
				fileCache[vma.Path] = file
			}
			return file, fileOffset, nil
		}
	}

	delete(vmaCache, pid)
	return nil, 0, fmt.Errorf("address 0x%x not found in process %d's memory map", vaddr, pid)
}

func parseProcMaps(pid uint32) ([]VMAEntry, error) {
	filePath := fmt.Sprintf("/proc/%d/maps", pid)
	content, err := os.ReadFile(filePath)
	if err != nil {
		return nil, err
	}

	var vmas []VMAEntry
	lines := strings.Split(string(content), "\n")
	for _, line := range lines {
		if line == "" {
			continue
		}

		fields := strings.Fields(line)
		if len(fields) < 6 {
			continue
		}

		addrRange := strings.Split(fields[0], "-")
		start, _ := strconv.ParseUint(addrRange[0], 16, 64)
		end, _ := strconv.ParseUint(addrRange[1], 16, 64)
		offset, _ := strconv.ParseUint(fields[2], 16, 64)

		path := fields[5]

		vmas = append(vmas, VMAEntry{
			StartAddr: start,
			EndAddr:   end,
			Offset:    offset,
			Path:      path,
		})
	}
	return vmas, nil
}